import cv2
import numpy as np

videoSrc = "../vtest.avi"

# 参数为0是打开摄像头，文件名是打开视频
cap = cv2.VideoCapture(videoSrc)

# 构造角点检测所需参数
feature_params = dict(maxCorners=100, qualityLevel=0.3, minDistance=7)

# lk_params = dict(winSize=(15, 15), maxLevel=2, criteria=(cv2.TERM_CRITERIA_EPS | cv2.TERM_CRITERIA_COUNT, 10, 0.03))
lk_params = dict(winSize=(15, 15), maxLevel=2)

color = np.random.randint(0, 255, (100, 3))

ret, old_frame = cap.read()

# 拿到第一帧图像并灰度化作为前一帧图片
old_gray = cv2.cvtColor(old_frame, cv2.COLOR_BGR2GRAY)
# old_gray = np.float32(old_gray)

# 返回所有检测特征点，需要输入图片，角点的最大数量，品质因子，
# minDistance=7如果这个角点里有比这个强的就不要这个弱的
p0 = cv2.goodFeaturesToTrack(old_gray, mask=None, **feature_params)

# print(p0)

# 创建一个mask, 用于进行横线的绘制
mask = np.zeros_like(old_frame)

while True:
    ret, frame = cap.read()
    frame_gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
    # 进行光流检测需要输入前一帧和当前图像及前一帧检测到的角点
    pl, st, err = cv2.calcOpticalFlowPyrLK(old_gray, frame_gray, p0, None, **lk_params)
    # 读取运动了的角点st == 1 表示检测到的运动物体，即v和u表示为0
    good_new = pl[st == 1]
    good_old = p0[st == 1]
    for i, (new, old) in enumerate(zip(good_new, good_old)):
        a, b = new.ravel()
        c, d = old.ravel()
        mask = cv2.line(mask, (a, b), (c, d), color[i].tolist(), 2)
        frame = cv2.circle(frame, (a, b), 5, color[i].tolist(), -1)
    # 将两个图片进行结合，并进行图片展示
    img = cv2.add(frame, mask)

    cv2.imshow("frame", img)
    k = cv2.waitKey(150) & 0xff
    if k == 27:
        break

    # 更新前一帧图片和角点的位置
    old_gray = frame_gray.copy()
    p0 = good_new.reshape(-1, 1, 2)
    # p0 = cv2.goodFeaturesToTrack(old_gray, mask=None, **feature_params)

cap.release()
cv2.destroyAllWindows()
